__attribute__((global, cooperative, reserve_r29_for_scavenge)) void kernel_upsample(int IC, int IH, int IW, int OH, int OW, float *__restrict dev_inp, float *__restrict dev_out)
{

    extern __shared__ float l2_buffer_inp[];

    // 获取线程的数量
    int thread_num = GetThreadNum();

    // 获取当前线程的全局索引
    int thread_idx = GetThreadIdx();

    // 获取线程的局部索引
    int local_idx = threadIdx.z * (blockDim.x * blockDim.y) + threadIdx.y * blockDim.x + threadIdx.x;

    // 定义与初始化DTE
    tops_dte_ctx_t dte_l1_inp;
    tops_dte_ctx_t dte_l1_out;

    __private_dte__ tops_dte_ctx_t dte_l2_inp;


    dte_l1_inp.init();
    dte_l1_out.init();
    dte_l2_inp.init();


    // 定义事件
    tops::event event_l1_inp;
    tops::event event_l1_out;

    tops::event event_l2_inp;


    // L3的形状
    int32_t global_inp_shape[3] = {IC, IH, IW};
    int32_t global_out_shape[3] = {IC, OH, OW};

    // l2的形状
    int32_t shared_inp_shape[3] = {1, IH, IW};

    // L1的形状
    int32_t private_inp_shape[3] = {1, IH, IW};
    int32_t private_out_shape[3] = {1, OH, OW};
    
    // mdspan绑定
    // L3
    tops::mdspan l3_input(tops::Global, dev_inp, global_inp_shape);
    tops::mdspan l3_output(tops::Global, dev_out, global_out_shape);

    // L2
    tops::mdspan l2_input(tops::Shared, l2_buffer_inp + (local_idx * IH * IW), shared_inp_shape);

    // L1
    int L1_BUFFER_SIZE = 1 * IH * IW;
    int L1_BUFFER_SIZE_OUT = 1 * OH * OW;

    __valigned__ float l1_buffer_inp[L1_BUFFER_SIZE];
    __valigned__ float l1_buffer_out[L1_BUFFER_SIZE_OUT];

    tops::mdspan l1_input(tops::Private, l1_buffer_inp, private_inp_shape);
    tops::mdspan l1_output(tops::Private, l1_buffer_out, private_out_shape);


    // 选取实际工作线程数
    int nr_working_threads = (IC <= thread_num) ? IC : thread_num;

    // 结束多余的工作线程
    if (thread_idx >= nr_working_threads)
    {
        return;
    }

    // 每个线程使用的通道数
    auto base_loops = IC / nr_working_threads;
    auto thread_ic_loops = base_loops + ((base_loops * nr_working_threads + thread_idx < IC) ? 1 : 0);
    
    // DTE绑定
    dte_l2_inp.connect(dte_l1_inp);

    for (int i = 0; i < thread_ic_loops; i++)
    {
        // ic维度的全局索引
        int global_ic_idx = i * nr_working_threads + thread_idx;

        event_l2_inp = tops::slice_async(dte_l2_inp, l2_input, l3_input, {global_ic_idx, 0, 0});
        event_l1_inp = tops::slice_async(dte_l1_inp, l1_input, l2_input, {0, 0, 0});

        tops::wait(event_l1_inp);

        if (i != 0)
        {
            tops::wait(event_l1_out);
        }

        for (int j = 0; j < OH; j++)
        {
            int global_oh_index = j;

            for (int k = 0; k < OW; k++)
            {
                int global_ow_index = k;

                // 对输入的位置进行放缩
                float ih = global_oh_index * (IH - 1) / (float)(OH - 1);
                float iw = global_ow_index * (IW - 1) / (float)(OW - 1);

                // 确定相邻的四个邻近像素
                int ih0 = floorf(ih);
                int iw0 = floorf(iw);
                int ih1 = ih0 + 1 < IH ? ih0 + 1 : ih0;
                int iw1 = iw0 + 1 < IW ? iw0 + 1 : iw0;

                // 计算小数部分
                float delta_h = ih - ih0;
                float delta_w = iw - iw0;

                // 计算权重
                float w1 = (1.0f - delta_h) * (1.0f - delta_w);
                float w2 = (1.0f - delta_h) * delta_w;
                float w3 = delta_h * (1.0f - delta_w);
                float w4 = delta_h * delta_w;

                
                float val1 = l1_buffer_inp[ih0 * IW + iw0];
                float val2 = l1_buffer_inp[ih0 * IW + iw1];
                float val3 = l1_buffer_inp[ih1 * IW + iw0];
                float val4 = l1_buffer_inp[ih1 * IW + iw1];

                float out_value = w1 * val1 + w2 * val2 + w3 * val3 + w4 * val4;

                l1_buffer_out[global_oh_index * OW + global_ow_index] = out_value;
            }
        }
        event_l1_out = tops::deslice_async(dte_l1_out, l3_output, l1_output, {global_ic_idx, 0, 0});

        if(i + 1 >= thread_ic_loops)
        {
            tops::wait(event_l1_out);
        }
    }

    // 销毁DTE
    dte_l1_inp.destroy();
    dte_l1_out.destroy();
    dte_l2_inp.destroy();
}
void GCU_UpSample(int IC, int IH, int IW, int OH, int OW, float *__restrict dev_inp, float *__restrict dev_out)
{
    static const size_t blocks = 2;
    static const size_t threads = 12;
    // 调用kernel
    kernel_upsample<<<blocks, threads, threads * IH * IW * sizeof(float)>>>(IC, IH, IW, OH, OW, dev_inp, dev_out);
}